import numpy as np
import cv2
import os

def img2vector(imgfilename):
    img = cv2.imread(imgfilename, cv2.IMREAD_GRAYSCALE)
    rows, columns = img.shape
    #print(img.shape)
    vec = img.reshape(rows * columns)
    return vec


def distance(vector1, vector2):
    diff = vector1 - vector2  # 差
    double_diff = diff ** 2  # 差平方
    sum_of_double_diff = double_diff.sum()  # 差平方和
    dist = np.sqrt(sum_of_double_diff)
    return dist


# training_image_matrix是训练集图像数据矩阵，class_vector是类别向量
def classify(img, training_image_matrix, class_vector, k):
    N = len(training_image_matrix)  # 训练集的图像个数
    dist = np.zeros((N))  # 距离向量初始化为全 0
    # 计算待识别图像与训练集各个图像之间的欧几里德距离
    for i in range(N):
        dist[i] = distance(img, training_image_matrix[i])
    # 对距离进行排序
    sorted_index = dist.argsort() # numpy的间接排序，返回的是元素的下标
    match_count = {} # 用字典方式，便于数据成对
    # K 近邻   拿到前 K 个数值
    for i in range(k):
        match_class = class_vector[sorted_index[i]]
        match_count[match_class] = match_count.get(match_class,0) + 1 # 找不到则缺省为0

    # 对字典进行排序 从高到底 降序
    match_count_in_order = sorted(match_count.items(),key=lambda item:item[1],reverse=True)
    decided = match_count_in_order[0][0]
    return decided


training_dir = "../training"
sub_dir_and_files = os.listdir(training_dir)

sub_dirs = []
# 如果是目录
for x in sub_dir_and_files:
    if os.path.isdir(training_dir + "/" + x):
        sub_dirs.append(x)

# 计算装备图像总数
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(training_dir + "/" + subdir))

# 初始化训练图像数据矩阵 (N 行，128*128列) 和装备向量（长度为N）
training_img_matrix = np.zeros((N, 128 * 128))  # 每个图像一行数据
training_equipment_vector = [''] * N

i = 0  # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(training_dir + "/" + subdir)
    for image in image_files:
        # 将图像转换为向量
        v = img2vector(training_dir + "/" + subdir + "/" + image)
        training_img_matrix[i] = v
        training_equipment_vector[i] = subdir
        i += 1


# 获取测试目录下的所有子目录
testing_dir = '../testing'
sub_dir_and_files = os.listdir(testing_dir)
sub_dirs = []
for x in sub_dir_and_files:
    if os.path.isdir(testing_dir + '/' + x):
        sub_dirs.append(x)


# 计算测试目录下所有装备图像的数量和
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(testing_dir + '/' + subdir))

# 初始化测试图像数据矩阵（N行，128*128列）和（相对应的）装备向量（长度为N）
testing_image_matrix = np.zeros((N, 128*128)) # 每个测试图像对应一行所有像素值
testing_equipment_vector = [""]*N # 与同下标位置测试图像相对应的装备名称


# 将三个子目录下所有图像文件读入图像数据矩阵并标识装备类别名称
i = 0 # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(testing_dir + "/" + subdir)
    for image in image_files:
        # 图像不在当前目录，因此需要加上目录前缀
        v = img2vector(testing_dir+"/"+subdir+"/"+image)
        testing_image_matrix[i] = v
        testing_equipment_vector[i] = subdir
        i += 1


correct_count = 0
for i in range(N):
    k = 4
    recognized = classify(testing_image_matrix[i], training_img_matrix, training_equipment_vector, k)
    if recognized == testing_equipment_vector[i]:
        correct_count += 1
print('%d个测试对象判定正确%d个，正确率为：%.1f%%' %(N, correct_count, correct_count/N*100))

